#ifndef LOCAL_MATRIX_HPP
#define LOCAL_MATRIX_HPP

#include <iostream>
#include "mem.h"
#include <time.h>
#include <math.h>
#include <ios>
#include <bitset>
#include <vector>
#include <cstring>
#include <stdlib.h>
#include "mem.cpp"

class LocalMatrix{
        LocalMatrix( const LocalMatrix & other );
        LocalMatrix & operator=( const LocalMatrix & other );

        float * data_;
        int n_;
        int m_;
public:
        void print(){
                float * p_data = data_;
                for( int i = 0; i < n_; ++i ){
                        for(int j = 0; j < m_; ++j, ++p_data ){
                                std::cout.width(6);
                                std::cout << std::right << *p_data;
                        }
                        std::cout << "\n";
                }
                std::cout << std::endl;
        }
        void init_zero(){
                memset(data_, 0, n_ * m_ * sizeof(float) );
        }
        void init(){
                srand(time(NULL));
                for( int i = 0, N = n_ * m_; i < N; ++i ){
                        //data_[i] = (float)((rand() * 1.0) / RAND_MAX);
                        data_[i] = i;
                }
        }

        //b is transposed
        static void mul_transposed( LocalMatrix & a, LocalMatrix & b, LocalMatrix & c ){
                int n = a.n_;
                int m = a.m_;
                int k = c.n_;
                //std::cout << "n,m,k: " << n << " " << m << " " << k << std::endl;
                int c_index = 0;
                float * a_row = a.data_;
                float * b_col_start = b.data_;
                for( int i = 0; i < n; ++i, a_row += m ){
                        float * b_col = b_col_start;
                        for( int j = 0; j < k; ++j, ++c_index, b_col += m ){
                                float tmp = c.data_[c_index];
                                for( int index = 0; index < m; ++index ){
                //			std::cout << a_row[index] << " * " << b_col[index] << " + ";
                                        tmp += a_row[index] * b_col[index];
                                }
                //		std::cout << std::endl;
                                c.data_[c_index] = tmp;
                        }
                }
        }

        static void blockTranspose(float * from_data, float * to_data, int n, int a, int b){
            float * p_data_to_start = to_data;
            float * p_data_from = from_data;
            for( int i = 0; i < a; ++i, ++p_data_to_start ){
                    float * p_data_to = p_data_to_start;
                    for(int j = 0; j < b; ++j, ++p_data_from, p_data_to += n ){
                            *p_data_to = *p_data_from;
                    }
            }
        }


        void transposeByBlocks( int from_n, int to_n, int from_m, int to_m, float * to_data ){
    //        std::cout << "(" << from_n << ", " << to_n << ", " << from_m << ", " << to_m << ")" << std::endl;
            static const int BLOCK_SIZE = 16;
            int block_n = to_n - from_n;
            int block_m = to_m - from_m;
            if ( block_n < BLOCK_SIZE && block_m < BLOCK_SIZE ){
                float * from_data = data_ + n_ * from_m + from_n;
                to_data = to_data + m_ * from_n + from_m;
                blockTranspose( from_data, to_data, n_, block_n, block_m );
            }else if ( block_n > block_m ){
                int half_n = (from_n + to_n) / 2;
                transposeByBlocks( from_n, half_n, from_m, to_m, to_data );
                transposeByBlocks( half_n, to_n, from_m, to_m, to_data );
            }else{
                int half_m = (from_m + to_m) / 2;
                transposeByBlocks( from_n, to_n, from_m, half_m, to_data );
                transposeByBlocks( from_n, to_n, half_m, to_m, to_data );
            }
        }

        void transposeByBlocks( LocalMatrix & to_matrix ){
            float * new_data = to_matrix.data_;
            transposeByBlocks( 0, n_, 0, m_, new_data );
        }
        void transposeByBlocks(){
            float * new_data = (float *) mem_malloc_aligned ( n_ * m_ * sizeof(float), 256 );
            transposeByBlocks( 0, n_, 0, m_, new_data );
            std::swap( data_, new_data );
            std::swap( m_, n_ );
            mem_free_aligned( new_data );
        }

        void transpose(){
                float * new_data = (float *) mem_malloc_aligned ( n_ * m_ * sizeof(float), 256 );
                blockTranspose( data_, new_data,n_, n_, m_ );
                /*
                float * p_data_to_start = new_data;
                float * p_data_from = data_;
                for( int i = 0; i < n_; ++i, ++p_data_to_start ){
                        float * p_data_to = p_data_to_start;
                        for(int j = 0; j < m_; ++j, ++p_data_from, p_data_to += n_ ){
                                *p_data_to = *p_data_from;
                        }
                }*/
                std::swap( data_, new_data );
                std::swap( m_, n_ );
                mem_free_aligned( new_data );
        }

        void transpose_inplace(){
                int size = m_ * n_ - 1;
                float t; // holds element to be replaced, eventually becomes next element to move
                int next; // location of 't' to be moved
                int cycleBegin; // holds start of cycle
                int i; // iterator
                std::vector<bool> b; // hash to mark moved elements
                b.assign( size, false );
                b[0] = b[size] = true;
                i = 1; // Note that A[0] and A[size-1] won't move
                while ( i < size )
                {
                        cycleBegin = i;
                        t = data_[i];
                        do
                        {
                                // Input matrix [r x c]
                                // Output matrix 1
                                // i_new = (i*r)%(N-1)
                                next = (i* n_)%size;
                                std::swap( data_[next], t );
                                b[i] = 1;
                                i = next;
                        }
                        while( i != cycleBegin );

                        // Get Next Move (what about querying random location?)
                        for ( i = 1; i < size && b[i]; i++ )
                                ;
                }
                std::swap( m_, n_ );
        }
        LocalMatrix(int n, int m) :n_(n),m_(m) { data_ = (float *) mem_malloc_aligned( n * m * sizeof(float), 256 ); }
        ~LocalMatrix(){ mem_free_aligned( data_ ); }
};


#endif // LOCAL_MATRIX_HPP
